/************************************************************************

							Heroes IV
					  Copyright 2000, The 3DO Company

	------------------------------------------------------------------
  							  compressed_filter.cpp

	$Header: /resource_editor/compressed_filter.cpp $

	$NoKeywords: $

 ************************************************************************/

// kill the 255+ character identifier warning.  Templates tend to exceed that easily.
#include "precompiled.h"
#pragma warning( disable: 4786 )

#include "compressed_filter.h"

#include <algorithm>
#include <assert.h>
#include <stdio.h>
#include <string.h>

#include "debug_switch.h"
#include "elements_of.h"
#include "game_time.h"
#include "zutil.h"

static int const k_buffer_size = 1024*8;
static int const k_buffer_size_2X  = k_buffer_size*2;

/* gzip flag byte */
#define ASCII_FLAG   0x01 /* bit 0 set: file probably ascii text */
#define HEAD_CRC     0x02 /* bit 1 set: header CRC present */
#define EXTRA_FIELD  0x04 /* bit 2 set: extra field present */
#define ORIG_NAME    0x08 /* bit 3 set: original file name present */
#define COMMENT      0x10 /* bit 4 set: file comment present */
#define RESERVED     0xE0 /* bits 5..7: reserved */

/////////////////////////////////////////////////////////////////////////////
// Unnamed namespace

namespace
{

/////////////////////////////////////////////////////////////////////////////
int const k_gz_magic[ 2 ] = { 0x1f, 0x8b };

} // Unnamed namespace

/////////////////////////////////////////////////////////////////////////////
//	t_deflate_filter members

/////////////////////////////////////////////////////////////////////////////
inline char * t_deflate_filter::get_input_buffer()
{
	return &m_buffer[ 0 ];
}

/////////////////////////////////////////////////////////////////////////////
inline char * t_deflate_filter::get_output_buffer()
{
	return &m_buffer[ k_buffer_size ];
}

/////////////////////////////////////////////////////////////////////////////
t_deflate_filter::t_deflate_filter( std::streambuf & destination, t_int32 datatype ) :
	m_dest( destination ),
	m_buffer( new char[k_buffer_size_2X] ),
	m_crc( crc32( 0L, Z_NULL, 0 ) ),
	m_open( false )
{
	// datatype determines type of compression and strategy used
	m_datatype = datatype;

	int level;
	int strategy;
	switch(m_datatype)
	{
	case k_raw_uncompressed_binary:
		level		= Z_NO_COMPRESSION;
		strategy	= Z_DEFAULT_STRATEGY;
		break;
	case k_low_compression_binary:
		level		= Z_BEST_SPEED;
		strategy	= Z_DEFAULT_STRATEGY;
		break;
	case k_medium_compression_binary:
		level		= Z_DEFAULT_COMPRESSION;
		strategy	= Z_DEFAULT_STRATEGY;
		break;
	case k_high_compression_binary:
		level		= Z_BEST_COMPRESSION;
		strategy	= Z_DEFAULT_STRATEGY;
		break;
	default:
		level		= Z_DEFAULT_COMPRESSION;
		strategy	= Z_DEFAULT_STRATEGY;
		break;
	}
	
	//k_raw_uncompressed_binary doesn't have any compression.

	// Initialize the streambuf pointers
	char * input_buffer = get_input_buffer();
	
	if (m_datatype == k_raw_uncompressed_binary)
	{
		setg( NULL, NULL, NULL );
		setp( input_buffer, input_buffer, input_buffer + k_buffer_size_2X );
	}
	else
	{
		setg( NULL, NULL, NULL );
		setp( input_buffer, input_buffer, input_buffer + k_buffer_size );

		// Open the deflate stream
		m_zstream.next_in = reinterpret_cast< Bytef * >( input_buffer );
		m_zstream.avail_in = 0;
		m_zstream.next_out = reinterpret_cast< Bytef * >( get_output_buffer() );
		m_zstream.avail_out = k_buffer_size;
		m_zstream.zalloc = Z_NULL;
		m_zstream.zfree = Z_NULL;
		int result = deflateInit2( &m_zstream, level, Z_DEFLATED, -MAX_WBITS, DEF_MEM_LEVEL, strategy );
		if ( result != Z_OK )
			throw t_create_failure();
		
		// Write out a simple gzip header.
		char gz_header[ 10 ];
		std::fill_n( gz_header, ELEMENTS_OF( gz_header ), '\0' );
		gz_header[ 0 ] = k_gz_magic[ 0 ];
		gz_header[ 1 ] = k_gz_magic[ 1 ];
		gz_header[ 2 ] = Z_DEFLATED;
		gz_header[ 9 ] = OS_CODE;
		if ( m_dest.sputn( gz_header, ELEMENTS_OF( gz_header ) ) < ELEMENTS_OF( gz_header ) )
			throw t_create_failure();
	}

	m_open = true;
}

/////////////////////////////////////////////////////////////////////////////
t_deflate_filter::~t_deflate_filter()
{
	try
	{
		close();
	}
	catch ( ... )
	{
	}
}

/////////////////////////////////////////////////////////////////////////////
int t_deflate_filter::close()
{
	if ( !m_open )
		return 0;
	m_open = false;

	if (m_datatype == k_raw_uncompressed_binary)
	{
		char * input_buffer  = get_input_buffer();

		t_int32 buffered_raw_bytes = pptr() - pbase();
			
		if ( m_dest.sputn( input_buffer, buffered_raw_bytes ) < buffered_raw_bytes )
			return -1;
	}
	else
	{
		char * input_buffer = get_input_buffer();
		char * output_buffer = get_output_buffer();
		
		assert( pptr() >= input_buffer && pptr() <= input_buffer + k_buffer_size );
		
		// Flush all pending output
		m_zstream.next_in = reinterpret_cast< Bytef * >( input_buffer );
		m_zstream.avail_in = pptr() - input_buffer;
		
		m_crc = crc32( m_crc, m_zstream.next_in, m_zstream.avail_in );
		
		int result;
		while ( ( result = deflate( &m_zstream, Z_FINISH ) ) == Z_OK || result == Z_BUF_ERROR )
		{
			assert( result == Z_OK || ( result == Z_BUF_ERROR && m_zstream.avail_out == 0 ) );
			if ( m_dest.sputn( output_buffer, k_buffer_size - m_zstream.avail_out ) < k_buffer_size - m_zstream.avail_out )
				return -1;
			m_zstream.next_out = reinterpret_cast< Bytef * >( output_buffer );
			m_zstream.avail_out = k_buffer_size;
		}
		assert( result == Z_STREAM_END );
		if ( m_dest.sputn( output_buffer, k_buffer_size - m_zstream.avail_out ) < k_buffer_size - m_zstream.avail_out )
			return -1;
		
		// Write out the crc and total bytes in
		if (	put( m_crc ) == -1
			||	put( m_zstream.total_in ) == -1 )
			return -1;
		
		// Close the deflate stream
		result = deflateEnd( &m_zstream );
		if ( result != Z_OK )
			return -1;
	}

	return 0;
}

/////////////////////////////////////////////////////////////////////////////
int t_deflate_filter::put( unsigned long x )
{
	if ( m_dest.sputc( static_cast< char >( x & 0xff ) ) == traits_type::eof() )
		return -1;
	if ( m_dest.sputc( static_cast< char >( ( x >> 8 ) & 0xff ) ) == traits_type::eof() )
		return -1;
	if ( m_dest.sputc( static_cast< char >( ( x >> 16 ) & 0xff ) ) == traits_type::eof() )
		return -1;
	if ( m_dest.sputc( static_cast< char >( ( x >> 24 ) & 0xff ) ) == traits_type::eof() )
		return -1;
	return 0;
}

/////////////////////////////////////////////////////////////////////////////
int t_deflate_filter::sync()
{
	if (m_datatype == k_raw_uncompressed_binary)
		return sync_raw();
	
	return sync_compressed();
}

int t_deflate_filter::sync_raw()
{
	//handle uncompressed raw data

	//write out what is in the buffer
	
	char * input_buffer  = get_input_buffer();
	
	t_int32 buffered_raw_bytes = pptr() - pbase();
	
	if ( m_dest.sputn( input_buffer, buffered_raw_bytes ) < buffered_raw_bytes )
		return -1;
	
	// Forward this message to the destination streambuf
	if ( m_dest.pubsync() == -1 )
		return -1;
	
	// Reset the streambuf input buffer
	setp( input_buffer, input_buffer, input_buffer + k_buffer_size_2X );
	
	return 0;
}

int t_deflate_filter::sync_compressed()
{
	//handle compressed data
	
	char * input_buffer		= get_input_buffer();
	char * output_buffer	= get_output_buffer();
	
	assert( pptr() >= input_buffer && pptr() <= input_buffer + k_buffer_size );
	
	// Flush all pending output
	m_zstream.next_in = reinterpret_cast< Bytef * >( input_buffer );
	m_zstream.avail_in = pptr() - input_buffer;
	
	m_crc = crc32( m_crc, m_zstream.next_in, m_zstream.avail_in );
	
	int result;
	while ( ( result = deflate( &m_zstream, Z_SYNC_FLUSH ), m_zstream.avail_in > 0 ) )
	{
		assert( result == Z_OK || ( result == Z_BUF_ERROR && m_zstream.avail_out == 0 ) );
		if ( m_dest.sputn( output_buffer, k_buffer_size - m_zstream.avail_out ) < k_buffer_size - m_zstream.avail_out )
			return -1;
		m_zstream.next_out = reinterpret_cast< Bytef * >( output_buffer );
		m_zstream.avail_out = k_buffer_size;
	}
	assert( result == Z_OK || ( result == Z_BUF_ERROR && m_zstream.avail_out == 0 ) );
	if ( m_dest.sputn( output_buffer, k_buffer_size - m_zstream.avail_out ) < k_buffer_size - m_zstream.avail_out )
		return -1;
	
	// Reset the zstream output buffer
	m_zstream.next_out = reinterpret_cast< Bytef * >( output_buffer );
	m_zstream.avail_out = k_buffer_size;
	
	// Forward this message to the destination streambuf
	if ( m_dest.pubsync() == -1 )
		return -1;
	
	// Reset the streambuf input buffer
	setp( input_buffer, input_buffer, input_buffer + k_buffer_size );
	
	return 0;
}

/////////////////////////////////////////////////////////////////////////////
std::streambuf::int_type t_deflate_filter::overflow( int_type c )
{
	if ( c != traits_type::eof() )
	{	
		if (m_datatype == k_raw_uncompressed_binary)
			return overflow_raw(c);
		
		return overflow_compress(c);
	}
	
	return traits_type::not_eof( c );
}

/////////////////////////////////////////////////////////////////////////////
std::streambuf::int_type t_deflate_filter::overflow_raw( int_type c )
{
	// Handle uncompressed raw data

	// Write out what is in the buffer
	
	char * input_buffer  = get_input_buffer();
	
	if ( c != traits_type::eof() )
	{
		if ( m_dest.sputn( input_buffer, k_buffer_size_2X ) < k_buffer_size_2X )
			return traits_type::eof();
	}
	
	// Reset the streambuf input buffer
	setp( input_buffer, input_buffer, input_buffer + k_buffer_size_2X );
	
	// Put the char in the buffer.
	*input_buffer = traits_type::to_char_type( c );
	
	pbump( 1 );
	
	return traits_type::not_eof( c );
}

/////////////////////////////////////////////////////////////////////////////
std::streambuf::int_type t_deflate_filter::overflow_compress( int_type c )
{	
	// Handle compressed data

	char * input_buffer  = get_input_buffer();
	char * output_buffer = get_output_buffer();
	
	assert( pptr() >= input_buffer && pptr() <= input_buffer + k_buffer_size );

	// Push all pending output through the deflate stream
	m_zstream.next_in = reinterpret_cast< Bytef * >( input_buffer );
	m_zstream.avail_in = pptr() - input_buffer;
	
	m_crc = crc32( m_crc, m_zstream.next_in, m_zstream.avail_in );
	
	int result;
	while ( ( result = deflate( &m_zstream, Z_NO_FLUSH ), m_zstream.avail_in > 0 ) )
	{
		assert( result == Z_OK || ( result == Z_BUF_ERROR && m_zstream.avail_out == 0 ) );
		if ( m_dest.sputn( output_buffer, k_buffer_size - m_zstream.avail_out ) < k_buffer_size - m_zstream.avail_out )
			return traits_type::eof();
		m_zstream.next_out = reinterpret_cast< Bytef * >( output_buffer );
		m_zstream.avail_out = k_buffer_size;
	}
	
	assert( result == Z_OK || ( result == Z_BUF_ERROR && m_zstream.avail_out == 0 ) );
	
	// Reset the streambuf input buffer
	setp( input_buffer, input_buffer, input_buffer + k_buffer_size );
	
	// Put the char in the buffer.
	*input_buffer = traits_type::to_char_type( c );
	
	pbump( 1 );
	
	return traits_type::not_eof( c );
}

/////////////////////////////////////////////////////////////////////////////
// t_inflate_filter::t_data_error members

/////////////////////////////////////////////////////////////////////////////
t_inflate_filter::t_data_error::t_data_error()
	:	std::runtime_error( DEBUG_SWITCH( "Inflate stream data error", "" ) )
{
}

/////////////////////////////////////////////////////////////////////////////
// t_inflate_filter members

/////////////////////////////////////////////////////////////////////////////
inline char * t_inflate_filter::get_input_buffer()
{
	return &m_buffer[ 0 ];
}

/////////////////////////////////////////////////////////////////////////////
inline char * t_inflate_filter::get_output_buffer()
{
	return &m_buffer[ k_buffer_size ];
}

/////////////////////////////////////////////////////////////////////////////
int t_inflate_filter::get_char()
{
	if ( m_zstream.avail_in == 0 )
	{
		if ( m_end_of_file )
			return std::streambuf::traits_type::eof();

		std::streamsize bytes_read = m_source.sgetn( get_input_buffer(), k_buffer_size );
		if ( bytes_read < k_buffer_size )
			m_end_of_file = true;

		m_zstream.next_in = reinterpret_cast< Bytef * >( get_input_buffer() );
		m_zstream.avail_in = bytes_read;

		if ( m_zstream.avail_in == 0 )
			return std::streambuf::traits_type::eof();
	}

	int result = std::streambuf::traits_type::to_int_type( static_cast< char >( *m_zstream.next_in++ ) );
	--m_zstream.avail_in;

	return result;
}

/////////////////////////////////////////////////////////////////////////////
void t_inflate_filter::put_back_char( char c )
{
	assert( m_zstream.next_in > reinterpret_cast< Bytef * >( get_input_buffer() ) );
	assert( static_cast< char >( *( m_zstream.next_in - 1 ) ) == c );
	--m_zstream.next_in;
	++m_zstream.avail_in;
}

/////////////////////////////////////////////////////////////////////////////
inline int t_inflate_filter::must_get_char()
{
	int c = get_char();
	if ( c == std::streambuf::traits_type::eof() )
		throw t_data_error();
	return c;
}

/////////////////////////////////////////////////////////////////////////////
inline unsigned long t_inflate_filter::must_get_long()
{
	unsigned long result = static_cast< unsigned long >( must_get_char() );
	result += static_cast< unsigned long >( must_get_char() ) << 8;
	result += static_cast< unsigned long >( must_get_char() ) << 16;
	result += static_cast< unsigned long >( must_get_char() ) << 24;
	return result;
}

/////////////////////////////////////////////////////////////////////////////
t_inflate_filter::t_inflate_filter( std::streambuf & source , t_int32 datatype) :
	m_source( source ),
	m_buffer(new char[k_buffer_size_2X]),
	m_crc( crc32( 0L, Z_NULL, 0 ) ),
	m_is_compressed( true ),
	m_end_of_file( false ),
	m_stream_is_open( false ),
	m_open( false )
{
	m_datatype = datatype;

	if (m_datatype == k_raw_uncompressed_binary)
	{
		// Initialize the streambuf pointers
		char * input_buffer = get_input_buffer();
		setg( input_buffer, input_buffer, input_buffer);
		setp( NULL, NULL, NULL );
	}
	else
	{	
		// Initialize the streambuf pointers
		char * output_buffer = get_output_buffer();
		setg( output_buffer, output_buffer, output_buffer );
		setp( NULL, NULL, NULL );
		
		// Initialize the inflate stream
		m_zstream.next_in = reinterpret_cast< Bytef * >( get_input_buffer() );
		m_zstream.avail_in = 0;
		m_zstream.next_out = reinterpret_cast< Bytef * >( output_buffer );
		m_zstream.avail_out = k_buffer_size;
		m_zstream.zalloc = Z_NULL;
		m_zstream.zfree = Z_NULL;
		
		// Try reading the gzip magic numbers
		try
		{
			int magic1 = get_char();
			if ( magic1 == std::streambuf::traits_type::eof() )
				throw false;
			try
			{
				if ( magic1 != k_gz_magic[ 0 ] )
					throw false;
				
				int magic2 = get_char();
				if ( magic2 == std::streambuf::traits_type::eof() )
					throw false;
				if ( magic2 != k_gz_magic[ 1 ] )
				{
					put_back_char( magic2 );
					throw false;
				}
			}
			catch ( bool )
			{
				put_back_char( magic1 );
				throw;
			}
		}
		catch ( bool )
		{
			// The magic numbers aren't there so assume the data is
			// uncompressed
			m_is_compressed = false;
		}
		
		if ( m_is_compressed )
		{
			// Read the gzip header
			if ( must_get_char() != Z_DEFLATED )
				throw t_data_error();
			
			int flags = must_get_char();
			if ( ( flags & RESERVED ) != 0 )
				throw t_data_error();
			
			int n;
			for ( n = 6; n > 0; --n )
				must_get_char();
			
			// Discard the extra field
			if ( ( flags & EXTRA_FIELD ) != 0 )
			{
				unsigned len = static_cast< unsigned >( must_get_char() );
				len += static_cast< unsigned >( must_get_char() ) << 8;
				while ( len-- > 0 )
					must_get_char();
			}
			
			// Skip the original file name
			if ( ( flags & ORIG_NAME ) != 0 )
				while ( must_get_char() != 0 );
				
			// Skip the comment
			if ( ( flags & COMMENT ) != 0 )
				while ( must_get_char() != 0 );
					
			// Skip the header crc
			if ( ( flags & HEAD_CRC ) != 0 )
			{
				must_get_char();
				must_get_char();
			}
			
			// Open the inflate stream
			int result = inflateInit2( &m_zstream, -MAX_WBITS );
			if ( result != Z_OK )
				throw t_data_error();

			m_stream_is_open = true;
		}
	}

	m_open = true;
}

/////////////////////////////////////////////////////////////////////////////
t_inflate_filter::~t_inflate_filter()
{
	try
	{
		close();
	}
	catch ( ... )
	{
	}
}

/////////////////////////////////////////////////////////////////////////////
int t_inflate_filter::close()
{
	if ( !m_open )
		return 0;
	m_open = false;

	if (m_datatype == k_raw_uncompressed_binary)	
	{
		if ( egptr() > gptr() )
			m_source.pubseekoff( gptr() - egptr(), std::ios_base::cur, std::ios_base::in );
	}
	else
	{
		// Seek backwards to compensate for the read ahead input
		if ( m_zstream.avail_in > 0 )
			m_source.pubseekoff( -static_cast< std::streambuf::off_type >( m_zstream.avail_in ), std::ios_base::cur, std::ios_base::in );
		
		if ( m_is_compressed )
		{
			// If the inflate stream is open, close it
			if ( m_stream_is_open )
			{
				int result = inflateEnd( &m_zstream );
				if ( result != Z_OK )
					return -1;
			}
		}
		else
		{
			if ( egptr() > gptr() )
				m_source.pubseekoff( gptr() - egptr(), std::ios_base::cur, std::ios_base::in );
		}
	}

	return 0;
}

/////////////////////////////////////////////////////////////////////////////
std::streambuf::int_type t_inflate_filter::underflow()
{	
	if (m_datatype == k_raw_uncompressed_binary)
		return underflow_raw();
	
	return underflow_compress();
}

std::streambuf::int_type t_inflate_filter::underflow_raw()
{
	// handle the raw data 
	char * input_buffer = get_input_buffer();
	
	std::streamsize bytes_read = m_source.sgetn( input_buffer, k_buffer_size_2X );
	
	if ( bytes_read < k_buffer_size_2X )
		m_end_of_file = true;
	
	// Reset the streambuf output buffer
	setg( input_buffer, input_buffer, input_buffer + bytes_read );
	
	if ( egptr() > eback())
		return traits_type::to_int_type( *gptr() );
	
	return traits_type::eof();
}

std::streambuf::int_type t_inflate_filter::underflow_compress()
{
	// Handle compressed data

	char * input_buffer  = get_input_buffer();
	char * output_buffer = get_output_buffer();
	
	assert( gptr() == egptr() );
	assert( m_zstream.next_out == reinterpret_cast< Bytef * >( output_buffer ) );
	assert( m_zstream.avail_out == k_buffer_size );
	
	// Iterate while there is still room in the output buffer and there
	// is more input available.
	while ( m_zstream.avail_out > 0 && ( m_zstream.avail_in > 0 || !m_end_of_file ) )
	{
		// If there's currently nothing in the zstream input buffer, read
		// some more.
		if ( m_zstream.avail_in == 0 )
		{
			std::streamsize bytes_read = m_source.sgetn( input_buffer, k_buffer_size );
			if ( bytes_read < k_buffer_size )
				m_end_of_file = true;
			
			// Reset the zstream input buffer state.
			m_zstream.next_in = reinterpret_cast< Bytef * >( input_buffer );
			m_zstream.avail_in = bytes_read;
		}
		
		// If there is available input...
		if ( m_zstream.avail_in > 0 )
		{
			if ( m_is_compressed )
			{
				if ( m_stream_is_open )
				{
					int result = inflate( &m_zstream, Z_SYNC_FLUSH );
					
					switch ( result )
					{
					case Z_DATA_ERROR:
						throw t_data_error();
					}
					
					m_crc = crc32( m_crc, m_zstream.next_out + m_zstream.avail_out - k_buffer_size, k_buffer_size - m_zstream.avail_out );
					
					if ( result == Z_STREAM_END )
					{
						result = inflateEnd( &m_zstream );
						assert( result == Z_OK );
						m_stream_is_open = false;
						
						// Eat the crc
						must_get_long();
						// Eat the total output
						must_get_long();
						
						m_zstream.avail_in = 0;
						m_end_of_file = true;
						
						break;
					}
					else
						assert( result == Z_OK );
				}
			}
			else
			{
				size_t bytes_to_copy = m_zstream.avail_in < m_zstream.avail_out ? m_zstream.avail_in : m_zstream.avail_out;
				memcpy( m_zstream.next_out, m_zstream.next_in, bytes_to_copy );
				m_zstream.next_in += bytes_to_copy;
				m_zstream.avail_in -= bytes_to_copy;
				m_zstream.next_out += bytes_to_copy;
				m_zstream.avail_out -= bytes_to_copy;
			}
		}
	}
	
	// Reset the streambuf output buffer
	setg( output_buffer, output_buffer, output_buffer + k_buffer_size - m_zstream.avail_out );
	
	// Reset the zstream output buffer
	m_zstream.next_out = reinterpret_cast< Bytef * >( output_buffer );
	m_zstream.avail_out = k_buffer_size;
	
	return egptr() > eback() ? traits_type::to_int_type( *gptr() ) : traits_type::eof();
}
